#pragma once

#include <vector>
#include <algorithm>
#include <iostream>
#include "../core/i_alloc.h"
#include "../geometry/i_util.h"

namespace idl
{
  template <typename ThreedDataType, typename DepthMapType>
  class ProjectXYZToDepthMap
  {
  public:
    typedef std::pair<int, int> PixelCoord;
    enum PixelType { MISS, BUFFERED, HIT };
    ProjectXYZToDepthMap();
    ProjectXYZToDepthMap(int width, int height, int z_buf_rad);
    virtual ~ProjectXYZToDepthMap(){ cleanup(); };
    bool init(int width, int height, int z_buffer_rad);
    int project(const ThreedDataType* const threeds,
                const int n,
                const ThreedDataType K[9],
                const ThreedDataType R[9],
                const ThreedDataType t[3],
                DepthMapType* depthmap,
                const DepthMapType val_threshold,
                const bool canonical_space = true, //set this value to TRUE if you want the output depth map values to be in the canonical i.e., [I|0] space
                const DepthMapType val_empty = (DepthMapType)0.0,
                const bool verbose = true);
  protected:
    void cleanup();
    bool allocMemory(int m, int n);
    void buildNeighborLut(const int r);
  protected:
    bool _is_initialized;
    int _width, _height, _z_buf_rad;
    ThreedDataType* _range;
    int* _state;
    std::vector<PixelCoord> _offset_lut;
    
    struct _Projection
    {
      ThreedDataType range;
      DepthMapType depth;
      int u, v;
      inline bool operator <(const _Projection& p) const { return range < p.range; };
    };
  };

  template <typename ThreedDataType, typename DepthMapType>
  ProjectXYZToDepthMap<ThreedDataType, DepthMapType>::ProjectXYZToDepthMap() : _width(0), _height(0), _z_buf_rad(0), _state(NULL), _range(NULL), _is_initialized(false){}

  template <typename ThreedDataType, typename DepthMapType>
  ProjectXYZToDepthMap<ThreedDataType, DepthMapType>::ProjectXYZToDepthMap(int width, int height, int z_buf_rad) : _width(width), _height(height), _z_buf_rad(z_buf_rad), _state(NULL), _range(NULL), _is_initialized(false)
  {
    _is_initialized = allocMemory(_height, _width);
    if (_is_initialized)
    {
      buildNeighborLut(z_buf_rad);
    }
  }

  template <typename ThreedDataType, typename DepthMapType>
  bool ProjectXYZToDepthMap<ThreedDataType, DepthMapType>::allocMemory(int m, int n)
  {
    if (n == _width && m == _height)
    {
      if (_state == NULL) { _state = i_alloc<int>(n*m); }
      if (_range == NULL) { _range = i_alloc<ThreedDataType>(n*m); }
    }
    else
    {
      cleanup();
      _state = i_alloc<int>(n*m);
      _range = i_alloc<ThreedDataType>(n*m);
    }
    if (_state == NULL || _range == NULL)
    {
      return false;
    }
    return true;
  }

  template <typename ThreedDataType, typename DepthMapType>
  void ProjectXYZToDepthMap<ThreedDataType, DepthMapType>::cleanup()
  {
    i_free<int>(_state);
    i_free<ThreedDataType>(_range);
    _offset_lut.clear();
    _width = _height = _z_buf_rad = 0;
    _is_initialized = false;
  }

  template <typename ThreedDataType, typename DepthMapType>
  void ProjectXYZToDepthMap<ThreedDataType, DepthMapType>::buildNeighborLut(const int r)
  {
    if (r == _z_buf_rad)
    {
      if ( !_offset_lut.empty() ) { return; }
    }
    _offset_lut.clear();
    _z_buf_rad = r;
    const int r2 = i_sqr(r);
    for (int dy = -r; dy <= r; ++dy)
    {
      for (int dx = -r; dx <= r; ++dx)
      {
        if (dy == 0 && dx == 0)
        {
          continue;
        }
        if ( i_sqr(dx) + i_sqr(dy) <= r2 )
        {
          _offset_lut.push_back(std::make_pair(dx, dy));
        }
      }
    }
  }

  template <typename ThreedDataType, typename DepthMapType>
  bool ProjectXYZToDepthMap<ThreedDataType, DepthMapType>::init(int width, int height, int z_buf_rad)
  {
    if (_is_initialized)
    {
      cleanup();
    }
    _width = width;
    _height = height;
    if (!_width || !_height)
    {
      return false;
    }
    _is_initialized = allocMemory(height, width);
    if (_is_initialized)
    {
      buildNeighborLut(z_buf_rad);
    }
    return _is_initialized;
  }

  template <typename ThreedDataType, typename DepthMapType>
  int ProjectXYZToDepthMap<ThreedDataType, DepthMapType>::project(const ThreedDataType* const threeds,
                                                                  const int n,
                                                                  const ThreedDataType K[9],
                                                                  const ThreedDataType R[9],
                                                                  const ThreedDataType t[3],
                                                                  DepthMapType* depthmap,
                                                                  const DepthMapType val_threshold,
                                                                  const bool canonical_space,
                                                                  const DepthMapType val_empty,
                                                                  const bool verbose)
  {
    int i, iu, iv, nu, nv, index_c, index_n, state;
    const ThreedDataType *X = threeds;
    ThreedDataType Xc[3], uvw[3], cop[3], principal_axis[3];
    _Projection proj;
    std::vector<_Projection> projections;
    projections.reserve(_width * _height);

    if (!_is_initialized)
    {
      if (verbose)
      {
        std::cerr << "The object has not been initialized." << std::endl;
      }
      return -1;
    }
    /*clear memory buffer*/
    for (i = 0; i < _width*_height; ++i)
    {
      _state[i] = MISS;
      _range[i] = (ThreedDataType)0.0;
      depthmap[i] = val_empty;
    }
  
    i_cop_from_R_t(R, t, cop);
    i_principal_axis_vector_from_R(R, principal_axis);

    for (i = 0; i<n; i++, X += 3)
    {
      if (!i_point_in_front_of_camera(X, cop, principal_axis))
      {
        continue;
      }
      /*project the point X to image plane*/
      i_project_through_extrinsic(R, t, X, Xc);
      i_project_through_intrinsic(K, Xc, uvw);
      i_homogeneous_unitize3(uvw);
      iu = i_round(uvw[0]);
      iv = i_round(uvw[1]);    
      /*check if the projected 2d point is within the image plane*/
      if (i_within_2d(iu, iv, 0, 0, _width, _height))
      {
        proj.range = i_sqrt(i_squaresum_diff3(X, cop));
        proj.depth = canonical_space ? (DepthMapType)Xc[2] : (DepthMapType)X[2];
        proj.u = iu;
        proj.v = iv;
        projections.push_back(proj);
      }
    }
    if (projections.empty())
    {
      return -1;
    }
    std::sort(projections.begin(), projections.end());
    typename std::vector<_Projection>::const_iterator cit_proj = projections.begin();
    
    while (cit_proj < projections.end())
    {
      index_c = cit_proj->v * _width + cit_proj->u;
      if (_state[index_c] == MISS)
      {
        state = HIT;
        _range[index_c] = cit_proj->range;
      }
      else if (_state[index_c] == HIT)
      {
        state = BUFFERED;
      }
      else
      {
        if (_state[index_c] == BUFFERED && cit_proj->range - _range[index_c] < val_threshold)
        {
          state = HIT;
          _range[index_c] = cit_proj->range;
        }
        else
        {
          state = BUFFERED;
        }
      }
      if (state == HIT)
      {
        _state[index_c] = HIT;
        depthmap[index_c] = cit_proj->depth;
        /*fill in the local circular depth buffer*/
        iu = cit_proj->u;
        iv = cit_proj->v;
        std::vector<PixelCoord>::const_iterator cit_offset = _offset_lut.begin();
        while (cit_offset < _offset_lut.end())
        {
          nu = iu + cit_offset->first;
          nv = iv + cit_offset->second;
          if (i_within_2d(nu, nv, 0, 0, _width, _height))
          {
            index_n = nv*_width + nu;
            if (_state[index_n] == MISS)
            {
              _state[index_n] = BUFFERED;
              _range[index_n] = cit_proj->range;
            }
          }
          cit_offset++;
        }
      }
      cit_proj++;
    }
    return 0;
  }

}/* namespace idl */
